import cv2
import numpy as np
import os
import time
from voc_write import GEN_Annotations

if not os.path.exists('./images'):
    os.system('mkdir images')
if not os.path.exists('./Annotations'):
    os.system('mkdir Annotations')

while True:
    sensor = input('请输入摄像头路径：')
    if len(sensor) == 1:
        sensor = int(sensor)
    elif sensor.find('udp') >= 0 or \
         sensor.find('tcp') >= 0 or \
         sensor.find('rtmp') >= 0 or\
         sensor.find('rtsp') >= 0 or\
         sensor.find('http') >= 0:
        pass
    else:
        continue
    cam = cv2.VideoCapture(sensor)
    if cam.isOpened():
        break
    else:
        cam.release()

while True:
    print('按 S 记录背景')
    while True:
        ret,img = cam.read()
        cv2.imshow('camera',img)
        key = cv2.waitKey(1)
        if ord('s') == key or ord('S') == key:
            background = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
            break

    labe = input('请输入标签名称：')
    print('按 S 记录图像数据')
    print('按 Q 重新选择背景和标签')
    while True:
        ret,img = cam.read()
        gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
        diff = cv2.absdiff(background,gray)
        dst = cv2.threshold(diff,30,255,cv2.THRESH_BINARY)
        kernel = np.ones((8,8),dtype=np.uint8)
        dst = cv2.morphologyEx(dst,cv2.MORPH_CLOSE,kernel)
        contours,_ = cv2.findContours(dst,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
        max_size = 0
        index = 0
        for i in range(len(contours)):
            if len(contours[i]) < 4:
                continue
            x,y,w,h = cv2.boundingRect(contours[i])
            s = w*h
            if s > max_size:
                max_size = s
                index = i
        if max_size > 0:
            x,y,w,h = cv2.boundingRect(contours[index])
            xmin = x
            ymin = y
            xmax = x+w
            ymax = y+h
            img0 = img.copy()
            cv2.rectangle(img,(xmin,ymin),(xmax,ymax),(0,255,0),1)
            cv2.putText(img,labe,(x,y-5),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,100,0),1)
            cv2.imshow('camera',img)
            key = cv2.waitKey(1)
            if ord('s') == key or ord('S') == key:
                filename="%Y-%m-%d-%H-%M-%S",time.localtime()
                anno= GEN_Annotations(filename+'.jpg')
                anno.set_size(img0.shape[1],img0.shape[0],img0.shape[2])
                anno.add_pic_attr(labe,xmin,ymin,xmax,ymax)
                anno.savefile('./Annotations/'+filename+".xml")
                cv2.imwrite('./images/'+filename+'.jpg',img0)
            elif ord('q') == key or ord('Q') == key:
                break
        else:
            cv2.imshow('camera',img)
            cv2.waitKey(1)